# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: ziyan

import logging
import subprocess
import sys
import tempfile

from absl import app, flags
import jinja2

from coin.base.datetime_util import iterate_date
from coin.support.feed_tool.feed_cache.app.coin2_fast_feed_runner import (load_feed_writer_config,
                                                                          load_subscribers_setting,
                                                                          get_mea_from_config,
                                                                          get_groups_from_setting,
                                                                          convert_mea_to_toml_setting,
                                                                          load_symbol_groups_setting)


registered_mea = ['Spot.Binance.v1', 'Futures.Binance.v1', 'Futures.Bybit.v3','Futures.Okex.v5-swap','Futures.Dydx.v3','Spot.Bithumb.v2','Spot.Gdax.v1','Spot.Bitbank.v1',
                  'Spot.Upbit.v1','Spot.Btcturk.v1 ','Spot.Okex.v5','Spot.Kucoin.v1','Spot.Kraken.v1','Spot.Korbit.v1','Futures.Prex.v1']
registered_mea_05cn = ['Spot.Okex.v5', 'Futures.Okex.v5-swap'] 
registered_base = ['BTC','ETH', 'XRP', 'ADA', 'DOGE', 'MATIC', 'SOL', 'DOT', 'LTC', 'AVAX', 'SHIB', 'TRX', 'UNI', 'LINK', 'ATOM', 'XMR', 'XLM', 'FIL', 'BCHN', 'APT']
registered_quote = ['USD','USDT','KRW','JPY']
FLAGS = flags.FLAGS

class NoJobError(Exception):
  pass


def generate_script_setting():
  machines = [
      # "feed-02.ap-northeast-2.aws",
      "feed-05.ap-northeast-1.aws",
      "feed-05.cn-hongkong.aliyun",
      # "feed-02.eu-west-1.aws",
      # "feed-02.us-east-1.aws",
      # "feed-01.ap-northeast-1.aws.huobi",
  ]
  config, config_file = load_feed_writer_config(machines)
  subscribers = load_subscribers_setting()
  symbol_groups_config = load_symbol_groups_setting()
  settings = {}
  for m in machines:
    if (FLAGS.machine is not None) and (m not in FLAGS.machine):
      continue
    settings[m] = {}
    meas = get_mea_from_config(m, config)
    for mea in meas:
      if (FLAGS.mea is not None) and (mea not in FLAGS.mea):
        continue
      if mea not in registered_mea:
        continue
      if m == "feed-05.cn-hongkong.aliyun" and mea not in registered_mea_05cn:
        continue
      settings[m][mea] = {}
      recipes = get_recipe_from_config(m, mea, config, subscribers)
      groups = get_groups_from_setting(m, mea, config)
      for recipe in recipes:
        if (FLAGS.recipe is not None) and (recipe not in FLAGS.recipe):
          continue
        settings[m][mea][recipe] = {}
        for group in groups:
          if (FLAGS.group is not None) and (group not in FLAGS.group):
            continue
          if group_have_registered_symbol(mea,group,symbol_groups_config):
            settings[m][mea][recipe][group] = { 
                "worker": FLAGS.workers,
                "raw": FLAGS.raw,
            }
  return settings


def get_recipe_from_config(machine, mea, config, subscribers):
  channels = config[machine]['feed']['exchanges'][mea]['channels']
  assert len(channels) > 0, "Only support subscribe by channel"
  toml_mea = convert_mea_to_toml_setting(mea)
  recipes = [subscribers[toml_mea]['default_recipe']]
  return recipes

def group_have_registered_symbol(mea,group,symbol_groups_config):
  if 'Binance' in mea:
    # binance needs almost all groups
    return True
  symbols = symbol_groups_config[convert_mea_to_toml_setting(mea)][group]
  for symbol in symbols:
    s = symbol.split('-')
    if "." not in s[1] and s[0] in registered_base and s[1] in registered_quote:
      return True
    elif  ".PERPETUAL" in s[1] and s[0] in registered_base and s[1].replace(".PERPETUAL","") in registered_quote:
      return True
  

def generate_slurm_script(start_time,
                          duration,
                          interval=None):
  
  settings = generate_script_setting()
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  template = template_env.get_template(
      'coin/support/feed_tool/feed_stats/app/feed_latency_stats/interval_stats_slurm_script.tmpl')
  return template.render(
                         start_time=start_time,
                         duration=duration,
                         settings=settings,
                         interval=interval)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_feed_stats_on_slurm(start_time,
                            duration,
                            dry_run=False,
                            sbatch_args=None,
                            verbose=False,
                            interval=None):
  try:
    script = generate_slurm_script(start_time,
                                   duration,
                                   interval=interval)
    
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  with tempfile.NamedTemporaryFile(mode='w') as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  return job_id


def main(argv):

  try:
    run_feed_stats_on_slurm(start_time=FLAGS.start_time,
                            duration=FLAGS.duration,
                            dry_run=FLAGS.script_only,
                            sbatch_args=FLAGS.sbatch_args,
                            verbose=True,
                            interval=FLAGS.interval)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string('start_time', None, 'yyyymmddTHHMM')
  
  flags.DEFINE_string('duration', 'PT30M', 'duration time')
  
  flags.DEFINE_bool('script_only', False,
                    'If true, it will print sbatch script and not run sbatch.')
  
  flags.DEFINE_string('sbatch_args',
                      '--wait --hint=nomultithread -n20 --mem-per-cpu=4G --job-name=coin2stats --priority=TOP',
                      'sbatch arguments')

  flags.DEFINE_list('machine', None, "Separate by comma")

  flags.DEFINE_list('mea', None, "Separate by comma")

  flags.DEFINE_list('recipe', None, "Separate by comma")

  flags.DEFINE_list('group', None, "Separate by comma")

  flags.DEFINE_list('workers', [0, 1, 2], "Separate by comma")
  
  flags.DEFINE_bool('raw', True, 'Run feed stats from raw feed.')

  flags.DEFINE_string('interval', '1M', 'interval time for interval feed generate, mmM or hhH')
  
  app.run(main)
